"""
任务管理API端点

提供任务查询和取消功能
"""

from typing import Optional
from fastapi import APIRouter, HTTPException, status, Query
from loguru import logger

from src.schemas.task import TaskStatusResponse, TaskStatusData, TaskCancelResponse, TaskCancelData, TaskStatus, TaskType
from src.schemas.base import success_response
from src.services.task_store import get_task_store


router = APIRouter(prefix="/task", tags=["任务管理"])


# 获取任务存储实例
task_store = get_task_store()


@router.get(
    "/{task_id}",
    response_model=TaskStatusResponse,
    summary="查询任务状态",
    description="""
    查询长时间运行任务的状态和结果。
    支持超时后继续查询完整结果。
    
    任务状态说明：
    - pending: 等待中
    - processing: 处理中
    - completed: 已完成
    - failed: 失败
    - timeout: 超时
    - cancelled: 已取消
    
    任务类型说明：
    - extraction: 字段提取任务
    - audit: 规则审核任务
    
    注意：如果提供了type参数，系统会验证任务类型是否匹配。
    如果任务存在但类型不匹配，将返回404错误。
    """
)
async def get_task_status(
    task_id: str,
    task_type: Optional[TaskType] = Query(default=None, alias="type", description="任务类型（extraction或audit），用于区分提取和审核任务"),
    include_details: bool = Query(default=True, description="是否包含详细结果")
) -> TaskStatusResponse:
    """
    查询任务状态
    
    Args:
        task_id: 任务ID
        task_type: 任务类型（可选），用于区分提取和审核任务
        include_details: 是否包含详细结果
        
    Returns:
        TaskStatusResponse: 任务状态响应
        
    Raises:
        HTTPException: 任务不存在或类型不匹配时抛出404错误
    """
    try:
        logger.info(f"查询任务状态: {task_id}, type={task_type}")
        
        # 根据任务类型添加前缀
        task = None
        internal_task_id = task_id
        
        if task_type == TaskType.EXTRACTION:
            # 明确指定提取任务
            internal_task_id = f"extraction_{task_id}"
            task = await task_store.get_task(internal_task_id)
        elif task_type == TaskType.AUDIT:
            # 明确指定审核任务
            internal_task_id = f"audit_{task_id}"
            task = await task_store.get_task(internal_task_id)
        else:
            # 未指定类型，先尝试提取任务，再尝试审核任务
            internal_task_id = f"extraction_{task_id}"
            task = await task_store.get_task(internal_task_id)
            if not task:
                internal_task_id = f"audit_{task_id}"
                task = await task_store.get_task(internal_task_id)
        
        if not task:
            logger.warning(f"任务不存在: {task_id} (内部ID: {internal_task_id})")
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"任务不存在: {task_id}"
            )
        
        # 任务类型验证（通过前缀已经保证了类型匹配）
        if task_type is not None and task.task_type != task_type:
            logger.warning(f"任务类型不匹配: {task_id}, 期望类型: {task_type}, 实际类型: {task.task_type}")
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"任务不存在或类型不匹配: {task_id}"
            )
        
        # 从 result 或 context 中提取 total_count 和 completed_count
        total_count = None
        completed_count = None
        
        if task.result and isinstance(task.result, dict):
            # 任务已完成，从 result 中获取完整统计信息
            if task.task_type == TaskType.EXTRACTION:
                # 提取任务：从 result 中获取 total_fields 和所有已处理的字段数
                total_count = task.result.get("total_fields")
                # completed_count = 所有已处理的字段数（成功 + 失败）
                success_count = task.result.get("success_count", 0)
                failed_count = task.result.get("failed_count", 0)
                completed_count = success_count + failed_count
            elif task.task_type == TaskType.AUDIT:
                # 审核任务：从 result.summary 中获取统计信息
                # 注意：passed/failed/pending 都算已完成，只要有审核结果就是完成
                summary = task.result.get("summary", {})
                if isinstance(summary, dict):
                    total_count = summary.get("total")
                    # completed_count = 所有有结果的审核规则数量（passed + failed + pending）
                    # 实际上就是 total，因为所有规则都已执行完成
                    passed = summary.get("passed", 0)
                    failed = summary.get("failed", 0)
                    pending = summary.get("pending", 0)
                    completed_count = passed + failed + pending
        else:
            # 任务处理中，从 context 中获取总数，根据进度计算已完成数
            if task.context and isinstance(task.context, dict):
                if task.task_type == TaskType.EXTRACTION:
                    # 提取任务：从 context 中获取 total_fields
                    total_count = task.context.get("total_fields")
                    if total_count is not None and task.progress is not None:
                        # 根据进度估算已完成数
                        completed_count = int(total_count * task.progress)
                elif task.task_type == TaskType.AUDIT:
                    # 审核任务：从 context 中获取 total_rules
                    total_count = task.context.get("total_rules")
                    if total_count is not None and task.progress is not None:
                        # 根据进度估算已完成数
                        completed_count = int(total_count * task.progress)
        
        # 处理result中的task_id（替换为原始ID）
        result_data = task.result
        if include_details and result_data and isinstance(result_data, dict):
            if result_data.get("task_id"):
                result_data = result_data.copy()
                result_data["task_id"] = task_id  # 替换为原始task_id
        
        # 构建响应数据（返回用户传入的原始task_id，不包含前缀）
        data = TaskStatusData(
            task_id=task_id,  # 返回原始task_id，不包含前缀
            task_type=task.task_type,
            status=task.status,
            progress=task.progress,
            result=result_data if include_details else None,
            total_count=total_count,
            completed_count=completed_count,
            error_message=task.error_message,
            created_at=task.created_at,
            updated_at=task.updated_at,
            completed_at=task.completed_at
        )
        
        logger.info(f"任务状态查询成功: {task_id}, 状态: {task.status}, 进度: {task.progress}")
        
        # 包装响应
        return success_response(
            data=data.model_dump(),
            message="success"
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"查询任务状态失败: {e}", exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"查询任务状态失败: {str(e)}"
        )


@router.post(
    "/{task_id}/cancel",
    response_model=TaskCancelResponse,
    summary="取消任务",
    description="""
    取消正在执行的任务。
    
    只有处于以下状态的任务可以取消：
    - pending: 等待中
    - processing: 处理中
    
    已完成、失败或已取消的任务无法再次取消。
    """
)
async def cancel_task(
    task_id: str
) -> TaskCancelResponse:
    """
    取消任务
    
    Args:
        task_id: 任务ID
        
    Returns:
        TaskCancelResponse: 取消响应
        
    Raises:
        HTTPException: 任务不存在或无法取消时抛出错误
    """
    try:
        logger.info(f"尝试取消任务: {task_id}")
        
        # 查询任务
        task = await task_store.get_task(task_id)
        
        if not task:
            logger.warning(f"任务不存在: {task_id}")
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"任务不存在: {task_id}"
            )
        
        # 检查任务状态
        if task.status not in [TaskStatus.PENDING, TaskStatus.PROCESSING]:
            logger.warning(f"任务状态为 {task.status}，无法取消: {task_id}")
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"任务状态为 {task.status.value}，无法取消"
            )
        
        # 取消任务
        success = await task_store.cancel_task(task_id)
        
        if not success:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail="取消任务失败"
            )
        
        logger.info(f"任务取消成功: {task_id}")
        
        # 构建响应数据
        data = TaskCancelData(
            task_id=task_id,
            status=TaskStatus.CANCELLED
        )
        
        # 包装响应
        return success_response(
            data=data.model_dump(),
            message="任务已取消"
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"取消任务失败: {e}", exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"取消任务失败: {str(e)}"
        )

